#include <atomic>
#include <mutex>
#include <memory>
#include <cstddef>
#include <new>
#include <set>
#include <iostream>

#include "memory.h"

namespace good_cptl
{

template<typename... Ts>
struct max_align_t_
{
  static constexpr std::size_t value()
  {
    return std::max({alignof(Ts)...});
  }
};

using max_align_v_ = max_align_t_<
  long double,
  double,
  float,
  long long int,
  long int,
  int,
  short int,
  bool,
  char,
  char16_t,
  char32_t,
  wchar_t,
  void*,
  std::max_align_t>;

constexpr std::size_t max_align_v = max_align_v_::value();
struct alignas(max_align_v) max_align_t {};

template<class Alloc>
class ThreadSafeArena
{
public:
  explicit ThreadSafeArena(Alloc alloc,
                           size_t min_block_size = 4096,
                           size_t align = alignof(std::max_align_t))
    : alloc_(alloc)
    , min_block_size_(min_block_size)
    , alignment_(align)
    , total_allocated_(0)
    , head_large_blocks_(nullptr)
    , head_blocks_(nullptr)
    , current_block_(nullptr)
  {}
  ~ThreadSafeArena()
  {
    while (head_large_blocks_) {
      head_large_blocks_->~LargeBlock();
      LargeBlock* next = head_large_blocks_->next;
      alloc_.deallocate((char*)head_large_blocks_, sizeof(LargeBlock));
      head_large_blocks_ = next;
    }
    while (head_blocks_) {
      head_blocks_->~Block();
      Block* next = head_blocks_->next;
      alloc_.deallocate((char*)head_blocks_, sizeof(Block));
      head_blocks_ = next;
    }
  }
  void* allocate(size_t size)
  {
    //std::cout << "size: " << size << std::endl;
    return allocate_aligned(size, alignment_);
  }

  void deallocate(void* ptr, size_t size) noexcept
  {
    std::lock_guard<std::mutex> lock(mutex_);
    const size_t aligned_size = align_up(size, alignment_);
    if (aligned_size > min_block_size_) {
      deallocate_large_block(ptr);
    } else {
      deallocate_block(ptr, aligned_size);
    }
    total_allocated_ -= size;
  }
  void* allocate_aligned(size_t size, size_t align)
  {
    std::lock_guard<std::mutex> lock(mutex_);
    if (align > alignment_) {
      align = alignment_;
    }
    const size_t aligned_size = align_up(size, align);
    //std::cout << "aligned_size: " << aligned_size << std::endl;
    if (aligned_size > min_block_size_) {
      total_allocated_ += aligned_size;
      return allocate_large_block(aligned_size);
    }
    if (try_reuse_blocks(aligned_size)) {
      //std::cout << "try_reuse_blocks, aligned_size: " << aligned_size << std::endl;
      void* ptr = current_block_->allocate(aligned_size, align);
      total_allocated_ += aligned_size;
      return ptr;
    }
    allocate_block(std::max(min_block_size_, aligned_size));
    void* ptr = current_block_->allocate(aligned_size, align);
    total_allocated_ += aligned_size;
    return ptr;
  }
public:
  void deallocate_block(void* ptr, size_t size)
  {
    auto cur_block = head_blocks_;
    while (cur_block) {
      if (ptr >= cur_block->ptr && ptr < cur_block->ptr + size) {
        cur_block->deallocate(size, alignment_);
        return;
      }
      cur_block = cur_block->next;
    }
  }
  void deallocate_large_block(void* ptr)
  {
    auto cur_large_block = head_large_blocks_;
    while (cur_large_block) {
      if (cur_large_block->ptr == ptr) {
        if (cur_large_block->prev) {
          cur_large_block->prev->next = cur_large_block->next;
        }
        if (cur_large_block->next) {
          cur_large_block->next->prev = cur_large_block->prev;
        }
        cur_large_block->~LargeBlock();
        alloc_.deallocate((char*)cur_large_block, sizeof(LargeBlock));
        cur_large_block = nullptr;
        return;
      }
      cur_large_block = cur_large_block->next;
    }
  }
  void* allocate_large_block(size_t size)
  {
    //std::cout << "allocate_large_block: " << size << std::endl;
    void* mem = alloc_.allocate(size);
    auto blk = new (mem) LargeBlock(size, alloc_);
    if (head_large_blocks_) {
      blk->next = head_large_blocks_;
      head_large_blocks_->prev = blk;
    }
    head_large_blocks_ = blk;
    return blk->ptr;
  }

  bool try_reuse_blocks(size_t size)
  {
    if (current_block_ && current_block_->remaining() >= size) {
      return true;
    }
    auto curr_block_ptr = head_blocks_;
    while (curr_block_ptr) {
      if (curr_block_ptr->remaining() >= size) {
        current_block_ = curr_block_ptr;
        return true;
      }
      curr_block_ptr = curr_block_ptr->next;
    }
    return false;
  }

  void allocate_block(size_t size)
  {
    size = align_up(size, alignment_);
    auto* new_block = new (alloc_.allocate(size)) Block(size, alloc_);
    new_block->next = head_blocks_;
    head_blocks_ = new_block;
    current_block_ = new_block;
  }
public:
  struct alignas(max_align_v) Block
  {
    Block(size_t s, Alloc& ac)
     : alloc(ac)
     , ptr(static_cast<char*>(alloc.allocate(s)))
     , next(nullptr)
     , size(s)
     , remaining_space(s)
    {}

    ~Block()
    {
      alloc.deallocate(ptr, size);
    }

    size_t remaining() const { return remaining_space; }
    void deallocate(size_t bytes, size_t align)
    {
      size_t aligned_size = align_up(bytes, align);
      remaining_space += aligned_size;
    }

    void* allocate(size_t bytes, size_t align)
    {
      char* start = ptr + (size - remaining());
      char* aligned_ptr = reinterpret_cast<char*>(align_up(reinterpret_cast<uintptr_t>(start), align));
      size_t actual_need = aligned_ptr - start + bytes;
      if (actual_need > remaining()) {
        return nullptr;
      }
      remaining_space -= actual_need;
      return aligned_ptr;
    }
    Alloc& alloc;
    char* ptr;
    Block* next;
    size_t size;
    size_t remaining_space;
  };
  struct alignas(max_align_v) LargeBlock
  {
    LargeBlock(size_t s, Alloc& ac)
     : alloc(ac)
     , ptr(static_cast<char*>(alloc.allocate(s)))
     , size(s)
     , prev(nullptr)
     , next(nullptr)
    {}

    ~LargeBlock()
    {
      alloc.deallocate(ptr, size);
    }

    Alloc& alloc;
    char* ptr;
    size_t size;
    LargeBlock* prev;
    LargeBlock* next;
  };
  static size_t align_up(size_t val, size_t align)
  {
    return (val + (align - 1)) & ~(align - 1);
  }
public:
  const size_t alignment_;
  const size_t min_block_size_;
  LargeBlock* head_large_blocks_;
  Block* head_blocks_;
  Block* current_block_;
  std::mutex mutex_;
  std::atomic<size_t> total_allocated_;
  Alloc alloc_;
};

class SysThreadSafeArena : public ThreadSafeArena<SysAllocator<char>>
{
public:
  explicit SysThreadSafeArena(size_t min_block_size = 4096,
                              size_t align = alignof(std::max_align_t))
    : ThreadSafeArena({}, min_block_size, align) {}
};

}